package gupao.concurrency.utils;

import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
public class ThreadPoolService {

    private final static ConcurrentHashMap<String, ThreadPoolExecutor> threadPoolMap = new ConcurrentHashMap<>();

    public static ThreadPoolExecutor getOrCreate(String bizName,
                                                 int corePoolSizeForCreateNew,
                                                 int maxPoolSizeForCreateNew
    ) {
        return getOrCreate(bizName, corePoolSizeForCreateNew,
                maxPoolSizeForCreateNew, new LinkedBlockingQueue<>(10000000), true, null);
    }

    public static ThreadPoolExecutor getOrCreate(String bizName,
                                                 int corePoolSizeForCreateNew,
                                                 int maxPoolSizeForCreateNew,
                                                 BlockingQueue<Runnable> taskQueue,
                                                 RejectedExecutionHandler handler
    ) {
        return getOrCreate(bizName, corePoolSizeForCreateNew,
                maxPoolSizeForCreateNew, taskQueue, true, handler);
    }

    public static ThreadPoolExecutor getOrCreate(String bizName,
                                                 int corePoolSizeForCreateNew,
                                                 int maxPoolSizeForCreateNew,
                                                 BlockingQueue<Runnable> taskQueueForCreateNew,
                                                 boolean isThreadDaemon,
                                                 RejectedExecutionHandler handler) {
        return threadPoolMap.computeIfAbsent(bizName, k -> create(bizName, corePoolSizeForCreateNew,
                maxPoolSizeForCreateNew, taskQueueForCreateNew, isThreadDaemon, handler));
    }

    public static void shutdown(String bizName) {
        ThreadPoolExecutor executor = threadPoolMap.get(bizName);
        if (executor != null) {
            executor.shutdown();
            log.info("thread pool {} shutdown.", bizName);
        }
    }

    public static List<Runnable> shutdownNow(String bizName) {
        ThreadPoolExecutor executor = threadPoolMap.get(bizName);
        if (executor != null) {
            log.info("thread pool {} shutdown now.", bizName);
            return executor.shutdownNow();
        }
        return new ArrayList<>();
    }

    private static ThreadPoolExecutor create(String prefix, int corePoolSize, int maxPoolSize,
                                             BlockingQueue<Runnable> queue,
                                             boolean isThreadDaemon,
                                             RejectedExecutionHandler handler) {
        return new ThreadPoolExecutor(corePoolSize, maxPoolSize, 30, TimeUnit.SECONDS, queue,
                new ThreadFactory() {
                    final AtomicInteger idx = new AtomicInteger(0);
                    @Override
                    public Thread newThread(Runnable r) {
                        Thread t = new Thread(r);
                        t.setName(prefix + "-" + idx.incrementAndGet());
                        t.setDaemon(isThreadDaemon);
                        t.setPriority(Thread.NORM_PRIORITY);
                        return t;
                    }
                }, handler == null ? new ThreadPoolExecutor.AbortPolicy() : handler);
    }
}
